package com.jtcoding.distributedlock.commons;

import com.google.common.collect.ImmutableList;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.core.RedisOperations;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.SessionCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.io.Serializable;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * spring boot redis的辅助类
 */
@Component
public class RedisHelper {

	@Autowired
	private RedisTemplate<String, Serializable> redisTemplate;
	@Autowired
	private StringRedisTemplate stringRedisTemplate;

	public Number rateLimit(String key, int limitCount, long limitPeriod) {
		String luaScript = buildRateLimitScript();
		RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
		Number count = redisTemplate.execute(redisScript, ImmutableList.of(key), limitCount, limitPeriod);
		return count;
	}

	/**
	 * setIfAbsent升级版，通过Lua的脚本，增加了超时时间
	 * 
	 * <p>
	 * 原redisTemplate.setIfAbsent成功之后断开连接，再执行redisTemplate.expire到达过期时间后key是依然存在redis中的
	 * </p>
	 * 
	 * @param key
	 * @param value
	 * @param seconds
	 * @return
	 */
	public boolean setIfAbsent(String key, String value, long seconds) {
		String luaScript = this.buildSetIfAbsentScript();
		RedisScript<Boolean> redisScript = new DefaultRedisScript<>(luaScript, Boolean.class);
		Object[] args = { value, seconds };
		return redisTemplate.execute(redisScript, ImmutableList.of(key), args);
	}

	public boolean delKey(String key) {
		String luaScript = this.buildDelScript();
		RedisScript<Boolean> redisScript = new DefaultRedisScript<>(luaScript, Boolean.class);
		return redisTemplate.execute(redisScript, ImmutableList.of(key));
	}

	/**
	 * setIfAbsent升级版，通过使用redis的事务Api中的multi()和exec()实现，增加了超时时间
	 * 
	 * <p>
	 * java 使用redis的事务时不能直接用Api中的multi()和exec()，这样multi()和exec()两次使用的stringRedisTemplate不是一个connect，会导致死锁
	 * </p>
	 * 
	 * @param key
	 * @param value
	 * @param seconds
	 * @return
	 */
	public boolean setIfAbsentMulti(String key, String value, long seconds) {
		SessionCallback<Boolean> sessionCallback = new SessionCallback<Boolean>() {
			List<Object> exec = null;

			@SuppressWarnings({ "unchecked", "rawtypes" })
			@Override
			public Boolean execute(RedisOperations operations) throws DataAccessException {
				operations.multi();
				stringRedisTemplate.opsForValue().setIfAbsent(key, value);
				stringRedisTemplate.expire(key, seconds, TimeUnit.SECONDS);
				exec = operations.exec();
				if (exec.size() > 0) {
					return (Boolean) exec.get(0);
				}
				return false;
			}
		};
		return stringRedisTemplate.execute(sessionCallback);
	}

	/**
	 * 构建限流的Lua脚本， 采用计数器限流
	 * <p>
	 * Note: lua脚本在redis执行可以保证操作的原子性, 减少网络延迟
	 * </p>
	 * <li>1、通过KEYS[1]获取传入的key参数</li>
	 * <li>2、通过ARGV[1]获取传入的limit参数</li>
	 * <li>3、通过redis.call方法，从缓存中get和key相关的值</li>
	 * <li>4、如果缓存中记录的数值不为nil，判断该数值是否会大于限制大小，如果超出表示该被限流，返回缓存值</li>
	 * <li>5、如果未超过，那么该key的缓存值+1，并返回缓存值+1</li>
	 * <li>6、如果是第一次调用，通过ARGV[2]获取传入的period参数设置该key的过期时间</li>
	 * 
	 * @return lua脚本
	 */
	public String buildRateLimitScript() {
		StringBuilder lua = new StringBuilder();
		lua.append("local c");
		lua.append("\nc = redis.call('get',KEYS[1])");
		// 调用不超过最大值，则直接返回
		lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
		lua.append("\nreturn c;");
		lua.append("\nend");
		// 执行计算器自加
		lua.append("\nc = redis.call('incr',KEYS[1])");
		lua.append("\nif tonumber(c) == 1 then");
		// 从第一次调用开始限流，设置对应键值的过期
		lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
		lua.append("\nend");
		lua.append("\nreturn c;");
		return lua.toString();
	}

	/**
	 * 构建setIfAbsent带超时时间的属性的Lua脚本,实现分布式锁
	 * <p>
	 * redisTemplete客户端，里面有一个setIfAbsent没有带超时时间的属性
	 * </p>
	 * <li>1、通过KEYS[1]获取传入的key参数</li>
	 * <li>2、通过ARGV[1]获取传入的value参数</li>
	 * <li>3、通过redis.call方法执行setnx判断key是否存在,若键 key 已经存在，则 setnx 命令不做任何动作。命令在设置成功时返回 1 ， 设置失败时返回 0</li>
	 * <li>4、如果执行成功，则获取ARGV[2]传入的expire，设置过期时间并且返回1</li>
	 * <li>5、如果执行失败，返回0</li>
	 * 
	 * @return lua脚本
	 */
	public String buildSetIfAbsentScript() {
		StringBuilder lua = new StringBuilder();
		lua.append("if 1 == redis.call('setnx',KEYS[1],ARGV[1]) then");
		lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
		lua.append("\nreturn 1;");
		lua.append("\nelse");
		lua.append("\nreturn 0;");
		lua.append("\nend;");
		return lua.toString();
	}

	public String buildDelScript() {
		StringBuilder lua = new StringBuilder();
		lua.append("if 1 == redis.call('EXISTS', KEYS[1]) then");
		lua.append("\nreturn redis.call('del', KEYS[1]);");
		lua.append("\nelse");
		lua.append("\nreturn 0;");
		lua.append("\nend;");
		return lua.toString();
	}
}
